package streammuxer

import (
	"fmt"
	"sync/atomic"

	"gitee.com/general252/gortc/pkg/rtprtcp"
	"gitee.com/general252/gortc/pkg/util"
)

var (
	_ StreamMuxer = (*StreamMux)(nil)
)

// StreamMux 如果一个mux处理能力不够, 使用n个mux, 均匀分片id到n个mux中 (例如: id%5=[0,1,2,3,4])
type StreamMux struct {
	streams *util.Map[StreamId, *Stream]
	ring    *rtprtcp.RingBuffer
	pullId  atomic.Uint64
}

// NewStreamMuxer 流路由, 所有的流都汇聚到这里, 然后从拉取
func NewStreamMuxer() StreamMuxer {
	ring, err := rtprtcp.NewRingBuffer(8192)
	if err != nil {
		panic(err)
	}

	s := &StreamMux{
		streams: util.NewMap[StreamId, *Stream](),
		ring:    ring,
	}

	s.run()

	return s
}

func (s *StreamMux) Close() {
	s.ring.Close()
}

func (s *StreamMux) run() {
	ring := s.ring

	go func() {
		for {
			v, ok := ring.Pull()
			if !ok {
				break
			}

			if data, ok := v.(*StreamData); ok {
				if stream, found := s.streams.Load(data.StreamId); found {
					stream.write(data)
				}
			}
		}
	}()
}

// AddStream 添加推流stream
func (s *StreamMux) AddStream(info StreamInfo) (*Stream, error) {
	if len(info.Id) == 0 {
		return nil, fmt.Errorf("no id")
	}

	stream, loaded := s.streams.LoadOrStore(info.Id, newStream(info))
	if loaded {
		return stream, fmt.Errorf("have exist")
	}

	return stream, nil
}

func (s *StreamMux) RemoteStream(stream *Stream) error {
	stream.close()
	s.streams.Delete(stream.GetInfo().Id)
	return nil
}

func (s *StreamMux) GetStream(id StreamId) (stream *Stream, ok bool) {
	if value, isOK := s.streams.Load(id); isOK {
		return value, true
	}

	return nil, false
}

func (s *StreamMux) RangeStream(f func(stream *Stream) bool) {
	s.streams.Range(func(key StreamId, value *Stream) bool {
		return f(value)
	})
}

// AddPull 添加拉流
func (s *StreamMux) AddPull(streamId StreamId, info PullInfo) (*Puller, error) {
	if info.OnPacket == nil {
		return nil, fmt.Errorf("error param")
	}

	stream, ok := s.streams.Load(streamId)
	if !ok {
		return nil, fmt.Errorf("not found")
	}

	id := s.pullId.Add(1)
	if id == 0 {
		id = s.pullId.Add(1)
	}

	puller := newPuller(info, streamId, pullId(id))

	_ = stream.AddPull(puller)

	return puller, nil
}

func (s *StreamMux) RemotePull(puller *Puller) error {
	stream, ok := s.streams.Load(puller.GetStreamId())
	if !ok {
		return fmt.Errorf("not found")
	}

	stream.RemotePull(puller)

	return nil
}

// Write 发送数据
func (s *StreamMux) Write(data *StreamData) error {
	s.ring.Push(data)

	return nil
}
